#include "stdafx.h"
#include "RWFile.h"
#include <vector>
#include <string>
#include <algorithm>
using namespace std;
#include <dbghelp.h>
#include <shlwapi.h>
#pragma comment(lib, "dbghelp")
#pragma comment(lib, "shlwapi")

string& trans_tolower(string& t)
{
	std::transform(t.begin(), t.end(), t.begin(), ::tolower);
	return t;
}

VOID GetFileDirList(string _FileOrDirName, vector<string>& _FileDirList, string _Root = "", string _Mask = "*.*")
{
	if (_Root == "")
	{
		_Root = _FileOrDirName;
	}
	WIN32_FIND_DATA _FindData;
	HANDLE hFind = INVALID_HANDLE_VALUE;
	if (_Mask == "*.*")
	{
		hFind = FindFirstFile((_FileOrDirName + "\\" + _Mask).c_str(), &_FindData);
		if (hFind == INVALID_HANDLE_VALUE)
			return;
		string _FindFileName;
		string _FullFileName;
		do {
			_FindFileName = _FindData.cFileName;
			_FullFileName = _FileOrDirName + "\\" + _FindFileName;
			if (_FindData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) {
				// 跳过当前目录和父目录
				if (_FindFileName == "." || _FindFileName == "..")
					continue;
				GetFileDirList(_FullFileName, _FileDirList, _Root, _Mask);
			}
			else {
				_FileDirList.push_back(_FullFileName);
			}
		} while (FindNextFile(hFind, &_FindData) != 0);
		FindClose(hFind);
	}
	else
	{
		hFind = FindFirstFile((_FileOrDirName + "\\" + _Mask).c_str(), &_FindData);
		if (hFind != INVALID_HANDLE_VALUE)
		{
			string _FindFileName;
			string _FullFileName;
			do {
				_FindFileName = _FindData.cFileName;
				_FullFileName = _FileOrDirName + "\\" + _FindFileName;
				if (_FindData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) {
				}
				else {
					_FileDirList.push_back(_FullFileName);
				}
			} while (FindNextFile(hFind, &_FindData) != 0);
			FindClose(hFind);
		}
		hFind = FindFirstFile((_FileOrDirName + "\\*.*").c_str(), &_FindData);
		if (hFind != INVALID_HANDLE_VALUE)
		{
			string _FindFileName;
			string _FullFileName;
			do {
				_FindFileName = _FindData.cFileName;
				_FullFileName = _FileOrDirName + "\\" + _FindFileName;
				if (_FindData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) {
					// 跳过当前目录和父目录
					if (_FindFileName == "." || _FindFileName == "..")
						continue;
					GetFileDirList(_FullFileName, _FileDirList, _Root, _Mask);
				}
			} while (FindNextFile(hFind, &_FindData) != 0);
			FindClose(hFind);
		}
	}
}

void DecPngList(vector<string>& _FileList)
{
	char pngHeader[] = { 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A };
	char pngXHeader[] = { 0x80, 0x70, 0x1F, 0x48, 0x1D, 0x2A, 0x1A, 0x3B };
	for (size_t i = 0; i < _FileList.size(); i++)
	{
		string& filename = _FileList[i];
		string kBuff;
		if (RWFile::LoadFile(filename, kBuff))
		{
			int len = kBuff.length();
			BYTE* data = (BYTE*)kBuff.c_str();
			if (memcmp(pngXHeader, data, 8))
			{
				continue;
			}
			memcpy(data, pngHeader, 8);
			for (int i = 8; i < len; i++)
			{
				data[i] ^= 0x6D;
			}
			RWFile::SaveFile(filename, kBuff);
		}
	}
}

void DecJpgList(vector<string>& _FileList)
{
	char jpgHeader[] = { 0xFF, 0xD8 };
	char jpgXHeader[] = { 0x1F, 0xF8 };
	for (size_t i = 0; i < _FileList.size(); i++)
	{
		string& filename = _FileList[i];
		string kBuff;
		if (RWFile::LoadFile(filename, kBuff))
		{
			int len = kBuff.length();
			if (len >= 5)
			{
				BYTE* data = (BYTE*)kBuff.c_str();
				if (memcmp(jpgXHeader, data, 2))
				{
					continue;
				}
				memcpy(data, jpgHeader, 2);
				for (int i = 2; i < len; i++)
				{
					data[i] ^= 0x6D;
				}
				RWFile::SaveFile(filename, kBuff);
			}
		}
	}
}

int main(int argc, char **argv)
{
	if (argc == 2)
	{
		string path = argv[1];
		if (!PathFileExists(path.c_str()))
		{
			printf("file not found!!!");
			return -1;
		}
		vector<string> _PngList;
		vector<string> _JpgList;
		if (PathIsDirectory(path.c_str()))
		{
			GetFileDirList(path, _PngList, "", "*.png");
			GetFileDirList(path, _JpgList, "", "*.jpg");
		}
		else
		{
			if (path.length() < 4)
			{
				printf("*.png, *.jpg not found!!!");
				return -1;
			}
			string ext = path.substr(path.length() - 4, 4);
			trans_tolower(ext);
			if (ext == ".png")
			{
				_PngList.push_back(path);
			}
			else if (ext == ".jpg")
			{
				_JpgList.push_back(path);
			}
			else
			{
				printf("*.png, *.jpg not found!!!");
				return -1;
			}
		}
		DecPngList(_PngList);
		DecJpgList(_JpgList);

	}
    return 0;
}